/*
    API Hooking Via Trampoline
    @5mukx
*/

use std::ptr::{null_mut, copy_nonoverlapping};
use windows::{
    core::{s, PCSTR, PCWSTR},
    Win32::{
        Foundation::HWND,
        System::{
            LibraryLoader::{GetModuleHandleA, GetProcAddress},
            Memory::{VirtualProtect, PAGE_EXECUTE_READWRITE, PAGE_PROTECTION_FLAGS},
        }, UI::WindowsAndMessaging::{MessageBoxA, MessageBoxW, MB_ICONINFORMATION, MB_OK},
    }
};

use widestring::WideCString;

const INTERCEPTOR_SIZE: usize = 14;

#[repr(C)]
struct ApiInterceptor {
    target_function: *mut std::ffi::c_void,
    replacement_function: *mut std::ffi::c_void,
    original_code: [u8; INTERCEPTOR_SIZE],
    original_protection: PAGE_PROTECTION_FLAGS,
}

impl ApiInterceptor {
    fn new() -> Self {
        ApiInterceptor {
            target_function: null_mut(),
            replacement_function: null_mut(),
            original_code: [0; INTERCEPTOR_SIZE],
            original_protection: PAGE_PROTECTION_FLAGS(0),
        }
    }
}

fn main(){

    unsafe {
        let user32 = GetModuleHandleA(s!("user32.dll")).expect("Failed to get user32.dll handle");
        
        let dialog_func = GetProcAddress(user32, s!("MessageBoxA")).expect("Failed to get MessageBox Address");

        let mut interceptor = ApiInterceptor::new();
        
        if !setup_interceptor(
            dialog_func as *mut std::ffi::c_void,
            custom_dialog as *mut std::ffi::c_void,
            &mut interceptor,
        ) {
            println!("[ERROR] Interceptor setup failed");
            return;
        }

        let text1 = s!("Testing 5mukx System");
        let caption1 = s!("System Info");

        MessageBoxA(
            None,
            text1,
            caption1,
            MB_OK | MB_ICONINFORMATION,
        );

        println!("[INFO] Activating API Interceptor...");

        if !activate_interceptor(&mut interceptor) {
            println!("[ERROR] Interceptor activation failed");
            return;
        }
        println!("[INFO] Interceptor activated");

         if !activate_interceptor(&mut interceptor) {
            println!("[ERROR] Interceptor activation failed");
            return;
        }
        println!("[INFO] Interceptor activated");

        let text2 = s!("Smukx Is Bad Guy...");
        let caption2 = s!("System Info");

        MessageBoxA(
            None,
            text2,
            caption2,
            MB_OK | MB_ICONINFORMATION,
        );

        println!("[INFO] Deactivating API interceptor...");
        if !deactivate_interceptor(&mut interceptor) {
            println!("[ERROR] Interceptor deactivation failed");
            return;
        }
        println!("[INFO] Interceptor deactivated");

        let text3 = s!("Smukx System Restored");
        let caption3 = s!("System Info");

        MessageBoxA(
            None,
            text3,
            caption3,
            MB_OK | MB_ICONINFORMATION,
        );
        
        println!("[INFO] PoC Demonstrated Successfully");
    }    

}

fn setup_interceptor(
    target_function: *mut std::ffi::c_void,
    replacement_function: *mut std::ffi::c_void,
    interceptor: &mut ApiInterceptor,
) -> bool{
    if target_function.is_null() || replacement_function.is_null(){
        return false;
    }

    interceptor.target_function = target_function;
    interceptor.replacement_function = replacement_function;

    unsafe {
        copy_nonoverlapping(
            target_function as *const u8,
            interceptor.original_code.as_mut_ptr(),
            INTERCEPTOR_SIZE,
        );

        let mut old_protection = PAGE_PROTECTION_FLAGS(0);
        if let Err(e) = VirtualProtect(
            target_function,
            INTERCEPTOR_SIZE,
            PAGE_EXECUTE_READWRITE,
            &mut old_protection,
            
        ) {
            println!("[!] Memory Protection change failed: {:?}", e);
            return false;
        }

        interceptor.original_protection = old_protection;
    }
    true
}

fn activate_interceptor(interceptor: &mut ApiInterceptor) -> bool {
    if interceptor.target_function.is_null() || interceptor.replacement_function.is_null() {
        return false;
    }

    unsafe {
        // far jump instruction (JMP)
        // 6 bytes for the JMP instruction (0xFF 0x25 0x00 0x00 0x00 0x00).
        // 8 bytes for the 64-bit address of the target function.

        let interceptor_code: [u8; INTERCEPTOR_SIZE] = [
            // JMP [RIP + 0]
            0xFF, 0x25, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];

        let patch = interceptor.replacement_function as u64;

        copy_nonoverlapping(
            &patch as *const _ as *const u8,
            (interceptor.target_function as *mut u8).offset(6),
            std::mem::size_of::<u64>(),
        );

        copy_nonoverlapping(
            interceptor_code.as_ptr(),
            interceptor.target_function as *mut u8,
            6,
        );
    }

    true
}

fn deactivate_interceptor(interceptor: &mut ApiInterceptor) -> bool {
    if interceptor.target_function.is_null() {
        return false;
    }

    unsafe {
        copy_nonoverlapping(
            interceptor.original_code.as_ptr(),
            interceptor.target_function as *mut u8,
            INTERCEPTOR_SIZE,
        );

        std::ptr::write_bytes(interceptor.original_code.as_mut_ptr(), 0, INTERCEPTOR_SIZE);

        let mut old_protection = PAGE_PROTECTION_FLAGS(0);
        if let Err(e) = VirtualProtect(
            interceptor.target_function,
            INTERCEPTOR_SIZE,
            interceptor.original_protection,
            &mut old_protection,
        ) {
            println!("[!] Memory protection restoration failed: {:?}", e);
            return false;
        }

        interceptor.target_function = null_mut();
        interceptor.replacement_function = null_mut();
        interceptor.original_protection = PAGE_PROTECTION_FLAGS(0);
    }

    true
}

unsafe extern "system" fn custom_dialog(
    hwnd: HWND,
    lp_text: PCSTR,
    lp_caption: PCSTR,
    u_type: u32,
) -> i32 {
    let text = unsafe { lp_text.to_string().unwrap_or_default() };
    let caption = unsafe { lp_caption.to_string().unwrap_or_default() } ;

    println!("[INFO] Dialog Parameters:");
    println!("\tText: {}", text);
    println!("\tCaption: {}", caption);

    let new_text = WideCString::from_str("5mukx Is a Good Guy").unwrap();
    let new_caption = WideCString::from_str("System Dialog").unwrap();

    unsafe {
        MessageBoxW(
            Some(hwnd),
            PCWSTR::from_raw(new_text.as_ptr()),
            PCWSTR::from_raw(new_caption.as_ptr()),
            windows::Win32::UI::WindowsAndMessaging::MESSAGEBOX_STYLE(u_type),
        ).0
    }
}